import base64
import inspect
from functools import wraps

from base.base_driver import DRIVER
from const import SCREEN_PATH


def exception_handling(obj):
    """
    异常处理装饰器，支持函数和类级别装饰
    对于类，会自动应用于所有以test_开头的方法
    """
    # 如果是函数装饰器
    if inspect.isfunction(obj):
        @wraps(obj)
        def wrapper(*args, **kwargs):
            DRIVER.start_recording_screen()
            try:
                return obj(*args, **kwargs)
            except Exception as e:
                result = DRIVER.stop_recording_screen()
                # 检查result类型，如果是字符串则转换为字节
                if isinstance(result, str):
                    result = result.encode('utf-8')
                with open(f"{SCREEN_PATH}/{obj.__name__}.mp4", "wb") as f:
                    f.write(result)
                raise e
        return wrapper
    
    # 如果是类装饰器
    elif inspect.isclass(obj):
        # 获取类的所有属性和方法
        for name, method in inspect.getmembers(obj, inspect.isfunction):
            # 只处理以test_开头的方法（测试方法）
            if name.startswith('test_'):
                # 装饰每个测试方法
                decorated_method = exception_handling(method)
                setattr(obj, name, decorated_method)
        return obj
    
    raise TypeError("exception_handling装饰器只能用于函数或类")
